#include "dma_funcs.hpp"
#include "simd_thl.h"
#include <sys/cdefs.h>
#include <tuple>
#include <utility>
#define __ldm_fix __attribute__ ((section (".ldm_fix")))
#define __ldm __attribute__ ((section (".ldm")))
#define __mpe__ __attribute__((sw_64_host))
#define __cpe__ __attribute__((sw_64_slave))
extern "C" {
extern void hthread_spawn(void *func, void *arg);
extern void hthread_join();
extern void hthread_init();
  
extern __ldm char  _CGN,_ROW,_COL,_PEN;
extern __ldm int   _MYID;

}
template<typename T>
void hthread_spawn(void (*func)(T*), T *arg){
  hthread_spawn((void*)func, arg);
}
template <typename G>
void __cpe__ hthread_run_lambda(G *g){
  char local_g[sizeof(G)];
  dma_get(g, &local_g[0], sizeof(G));
  (*(G*)local_g)();
  //(*g)();
}

template <typename G>
__always_inline __mpe__ void hparallel(G g){
  hthread_spawn((void*)__builtin_slave_func((void*)(void(*)(G*))hthread_run_lambda), &g);
  hthread_join();
}
enum intent { COPY_IN, COPY_OUT, FILL_0 };
template <typename T, int N>
struct harray1d {
  alignas(32) T data[N];
  T &operator[](int i){
    return data[i];
  }

  const T &operator[](int i) const {
    return data[i];
  }
  
  T *operator+(int i){
    return data + i;
  }
};

template<typename T, int N>
struct capture_span : harray1d<T, N>{
  capture_span(T *ptr, int n = N) {
    memcpy(this->data, ptr, N * sizeof(T));
  }
};

enum MemPtrMode {
  MP_IN,
  MP_OUT,
  MP_INOUT,
  MP_ZEROOUT
};

template<typename DT, MemPtrMode M>
struct hdirective{
  static constexpr MemPtrMode mode = M;
  typedef DT data_type;
  DT *&ptr;
};

template<typename DT, MemPtrMode M>
struct hdirective_restrict {
  static constexpr MemPtrMode mode = M;
  typedef DT data_type;
  DT * __restrict__ &ptr;
};

template<typename T, int N, MemPtrMode M>
struct hlinear_buffer {
  alignas(32) T buffer[N];
  //T * const mem;
  // __always_inline hlinear_buffer(T * __restrict__ &ptr_in) : mem(ptr_in) {
  // };
  // __always_inline hlinear_buffer(hdirective_restrict<T, M> dir) : mem(dir.ptr){
  // }

  // __always_inline hlinear_buffer(hdirective<T, M> dir) : mem(dir.ptr){
    
  // }
  // __always_inline void enter(off_t i, int count = N){
  //   if (M == MP_IN || M == MP_INOUT) {
  //     dma_getn(mem + i, buffer, count);
  //   }
  //   if (M == MP_ZEROOUT) {
  //     memset(buffer, 0, sizeof(T) * count);
  //   }
  // }
  // __always_inline void leave(off_t i, int count = N){
  //   if (M == MP_ZEROOUT || M == MP_INOUT || M == MP_OUT){
  //     dma_putn(mem + i, buffer, count);
  //   }
  // }
  __always_inline void enter(T *mem, off_t i, int count = N){
    if (M == MP_IN || M == MP_INOUT) {
      dma_getn(mem + i, buffer, count);
    }
    if (M == MP_ZEROOUT) {
      memset(buffer, 0, sizeof(T) * count);
    }
  }
  __always_inline void leave(T *mem, off_t i, int count = N){
    if (M == MP_ZEROOUT || M == MP_INOUT || M == MP_OUT){
      dma_putn(mem + i, buffer, count);
    }
  }

};

template<typename DataType>
__always_inline hdirective<DataType, MP_IN> hlinear_in(DataType *&ptr){
  return hdirective<DataType, MP_IN>{ptr};
}
template<typename DataType>
__always_inline hdirective<DataType, MP_OUT> hlinear_out(DataType *&ptr){
  return hdirective<DataType, MP_OUT>{ptr};
}

template<typename DataType>
__always_inline hdirective<DataType, MP_INOUT> hlinear_inout(DataType *&ptr){
  return hdirective<DataType, MP_INOUT>{ptr};
}

template<typename DataType>
__always_inline hdirective<DataType, MP_ZEROOUT> hlinear_zeroout(DataType *&ptr){
  return hdirective<DataType, MP_ZEROOUT>{ptr};
}

template<typename DataType>
__always_inline hdirective_restrict<DataType, MP_IN> hlinear_in(DataType * __restrict__ &ptr){
  return hdirective_restrict<DataType, MP_IN>{ptr};
}
template<typename DataType>
__always_inline hdirective_restrict<DataType, MP_OUT> hlinear_out(DataType * __restrict__ &ptr){
  return hdirective_restrict<DataType, MP_OUT>{ptr};
}

template<typename DataType>
__always_inline hdirective_restrict<DataType, MP_INOUT> hlinear_inout(DataType * __restrict__ &ptr){
  return hdirective_restrict<DataType, MP_INOUT>{ptr};
}

template<typename DataType>
__always_inline hdirective_restrict<DataType, MP_ZEROOUT> hlinear_zeroout(DataType * __restrict__ &ptr){
  return hdirective_restrict<DataType, MP_ZEROOUT>{ptr};
}

static long rpcc(){
  long ret;
  #ifdef __sw_slave__
  asm volatile("rcsr %0, 4\n\t": "=r"(ret) : : "memory");
  #endif
  return ret;
}

#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-value"
template<typename F, int STRIDE, int... Is>
__always_inline void _run_unroll(F &f, int ist, std::integer_sequence<int, Is...> seq){
  (int[]){(f(ist+Is * STRIDE), 0)...};

}

template<int BLKSIZE, int IDXSTRIDE, int BLKSTRIDE, typename ...Ts, typename F, int ...Is>
__always_inline void ___hiterate(int start, int stop, std::tuple<Ts...> dirs, F &&f, std::integer_sequence<int, Is...>){
  constexpr int unroll_factor = 4;
  std::tuple<std::remove_reference_t<decltype(std::get<Is>(dirs).ptr)>...> mem_save(std::get<Is>(dirs).ptr...);
  for (int ist = start; ist < stop; ist += BLKSTRIDE){
    auto buffers = std::make_tuple(hlinear_buffer<typename Ts::data_type, BLKSIZE, Ts::mode>{}...);
    (int[]){(std::get<Is>(buffers).enter(std::get<Is>(mem_save), ist), 0)...};
    std::tie(std::get<Is>(dirs).ptr...) = std::make_tuple((std::get<Is>(buffers).buffer - ist)...);
    for (int i = 0; i < BLKSIZE; i += IDXSTRIDE*unroll_factor){
      //__builtin_force_unroll(unroll_factor);
      //f(ist + i);
      _run_unroll<F, IDXSTRIDE>(f, ist + i, std::make_integer_sequence<int, unroll_factor>{});
    }
    (int[]){(std::get<Is>(buffers).leave(std::get<Is>(mem_save), ist), 0)...};
  }
  std::tie(std::get<Is>(dirs).ptr...) = mem_save;
}
#pragma GCC diagnostic pop

template<int BLKSIZE, int IDXSTRIDE, int BLKSTRIDE, typename ...Ts, int ...Is>
__always_inline void __hiterate(int start, int stop, std::tuple<Ts ...> args, std::integer_sequence<int, Is...> seq){
  ___hiterate<BLKSIZE, IDXSTRIDE, BLKSTRIDE>(start, stop, std::make_tuple(std::get<Is>(args)...), std::get<sizeof...(Ts)-1>(args), seq);
}

template<int BLKSIZE, int IDXSTRIDE, int BLKSTRIDE, typename ...Ts>
__always_inline void _hiterate(int start, int stop, Ts ... args){
  __hiterate<BLKSIZE, IDXSTRIDE, BLKSTRIDE>(start, stop, std::make_tuple(args...), std::make_integer_sequence<int, sizeof...(Ts) - 1>{});
}

// template<int BLKSIZE, int IDXSTRIDE=1, int BLKSTRIDE=BLKSIZE*64, typename ...Ts>
// auto hdistribute_iterate(int start, int stop, Ts ...args){
//   return hlinear_iterator<BLKSIZE, IDXSTRIDE, BLKSTRIDE, Ts...>(start, stop, args...);//.run(f);
// }
#define __sw_thl__
#ifdef __sw_thl__
//__v4df load_v4df(double *dat) { return __builtin_sw_vldd_f();}
template<typename DT>
__cpe__ __always_inline void putc_obj(const DT &val, int dest){
  constexpr int ALIGN = alignof(DT);
  constexpr int SIZE = sizeof(DT);

  static_assert(ALIGN >= 4);
  if (ALIGN == 4){
    const int *ptr = (int*)&val;
    for (int i = 0; i < SIZE; i += 4){
      asm volatile("putc %0, %1\n\t" :: "r"(ptr[i >> 2]), "r"(dest));
    }
  } else if (ALIGN == 8 || ALIGN == 16){
    const long *ptr = (long*)&val;
    for (int i = 0; i < SIZE; i += 8){
      asm volatile("putc %0, %1\n\t" :: "r"(ptr[i >> 3]), "r"(dest));
    }
  } else if (ALIGN >= 32){
    const __v4di *ptr = (__v4di*)&val;
    for (int i = 0; i < SIZE; i += 32) {
      asm volatile("putc %0, %1\n\t" :: "r"(ptr[i >> 5]), "r"(dest));
    }
  }
}

template<typename DT>
__cpe__ __always_inline DT getc_obj(){
  constexpr int ALIGN = alignof(DT);
  constexpr int SIZE = sizeof(DT);
  static_assert(ALIGN >= 4);
  DT val;
  if (ALIGN == 4){
    int *ptr = (int*)&val;
    for (int i = 0; i < SIZE; i += 4){
      __v8si tmp;
      tmp = __builtin_sw_slave_getcw(tmp);
      ptr[i >> 2] = __builtin_cast_v8si_si(tmp);
    }
  } else if (ALIGN == 8 || ALIGN == 16){
    long *ptr = (long*)&val;
    for (int i = 0; i < SIZE; i += 8){
      __v4di tmp;
      tmp = __builtin_sw_slave_getco(tmp);
      ptr[i >> 3] = __builtin_cast_v4di_di(tmp);
    }
  } else if (ALIGN >= 32){
    __v4di *ptr = (__v4di*)&val;
    for (int i = 0; i < SIZE; i += 32) {
      __v4di tmp;
      tmp = __builtin_sw_slave_getco(tmp);
      ptr[i >> 5] = tmp;
    }
  }
  return val;
}

template<typename DT>
__cpe__ __always_inline void putr_obj(const DT &val, int dest){
  constexpr int ALIGN = alignof(DT);
  constexpr int SIZE = sizeof(DT);

  static_assert(ALIGN >= 4);
  if (ALIGN == 4){
    const int *ptr = (int*)&val;
    for (int i = 0; i < SIZE; i += 4){
      //警告，__builtin_sw_slave_putxx有bug，会被优化掉，需要等啥时候打个补丁
      asm volatile("putr %0, %1\n\t" :: "r"(ptr[i >> 2]), "r"(dest));
    }
  } else if (ALIGN == 8 || ALIGN == 16){
    const long *ptr = (long*)&val;
    for (int i = 0; i < SIZE; i += 8){
      //__builtin_sw_slave_putro(__builtin_cast_di_v4di(ptr[i>>3]), dest);
      asm volatile("putr %0, %1\n\t" :: "r"(ptr[i >> 3]), "r"(dest));
    }
  } else if (ALIGN >= 32){
    const __v4di *ptr = (__v4di*)&val;
    for (int i = 0; i < SIZE; i += 32) {
      asm volatile("putr %0, %1\n\t" :: "r"(ptr[i >> 5]), "r"(dest));
    }
  }
}

template<typename DT>
__cpe__ __always_inline DT getr_obj(){
  constexpr int ALIGN = alignof(DT);
  constexpr int SIZE = sizeof(DT);
  static_assert(ALIGN >= 4);
  DT val;
  if (ALIGN == 4){
    int *ptr = (int*)&val;
    for (int i = 0; i < SIZE; i += 4){
      __v8si tmp;
      tmp = __builtin_sw_slave_getrw(tmp);
      ptr[i >> 2] = __builtin_cast_v8si_si(tmp);
    }
  } else if (ALIGN == 8 || ALIGN == 16){
    long *ptr = (long*)&val;
    for (int i = 0; i < SIZE; i += 8){
      __v4di tmp;
      tmp = __builtin_sw_slave_getro(tmp);
      ptr[i >> 3] = __builtin_cast_v4di_di(tmp);
    }
  } else if (ALIGN >= 32){
    __v4di *ptr = (__v4di*)&val;
    for (int i = 0; i < SIZE; i += 32) {
      __v4di tmp;
      tmp = __builtin_sw_slave_getro(tmp);
      ptr[i >> 5] = tmp;
    }
  }
  return val;
}

#endif

template<typename DT, typename AT>
__always_inline __cpe__ DT hreduce(DT in, AT &&add){
  asm volatile("sync %0\n\tsynr %0\n\t" ::"r"(0xff) : "memory");
  DT ret = in;
  for (int i = 1; i < 8; i += i){
    asm volatile("synr %0\n\t" :: "r"(0xff) : "memory");
    if ((_COL & i - 1) == 0) {
      if (_COL & i) {
        putr_obj(ret, _COL ^ i);
      } else {
        ret = add(ret, getr_obj<DT>());
      }
    }
  }
  asm volatile("synr %0\n\t" :: "r"(0xff) : "memory");
  for (int i = 1; i < 8; i += i){
    asm volatile("synr %0\n\t" :: "r"(0xff) : "memory");       
    if ((_ROW & i - 1) == 0) {
      if (_ROW & i) {
        putc_obj(ret, _ROW ^ i);
      } else {
        ret = add(ret, getc_obj<DT>());
      }
    }
    asm volatile("synr %0\n\t" :: "r"(0xff) : "memory");        
  }
  return ret;
}

template<typename DT>
__always_inline __cpe__ DT hbcast(DT in, int root = 0){
  DT val = in;
  asm volatile("sync %0\n\t" :: "r"(0xff));
  int rootrow = root >> 3;
  if (_MYID == root){
    putr_obj(val, 8);
    putc_obj(val, 8);
  } else if (_ROW == rootrow){
    val = getr_obj<DT>();
    putc_obj(val, 8);
  } else {
    val = getc_obj<DT>();
  }
  return val;
}
template<typename DT, typename AT>
__always_inline __cpe__ DT hallreduce(DT in, AT &&add) {
  DT ret = hreduce(in, add);
  asm volatile("sync %0\n\tsynr %0\n\t" ::"r"(0xff) : "memory");  
  return hbcast(ret, 0);
}
